import pandas as pd
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
from sklearn.metrics import silhouette_score

df=pd.read_csv('drink.txt')

k=[2,3,4,5,6]
#误差平方和
sse=[]
#轮廓系数
si=[]
for i in k:
    model=KMeans(n_clusters=i)
    model.fit(df)
    #存入误差平方和
    sse.append(model.inertia_)
    #获取预测所得样本类别
    label=model.predict(df)
    #存入轮廓系数值
    si.append(silhouette_score(df,label))
#画出肘部法则图
plt.plot(k,sse)
plt.show()
#轮廓系数图
plt.plot(k,si)
plt.show()